import os
import time

import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.decomposition import MiniBatchDictionaryLearning, DictionaryLearning
from sklearn.metrics import accuracy_score


# 加载数据集
def load_dataset(data_dir):
    X = []
    y = []
    labels = os.listdir(data_dir)
    label_to_idx = {label: idx for idx, label in enumerate(labels)}

    for label in labels:
        label_dir = os.path.join(data_dir, label)
        for img_file in os.listdir(label_dir):
            img_path = os.path.join(label_dir, img_file)
            img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)  # 如果是RGB图像，可以考虑使用 cv2.IMREAD_COLOR
            img = cv2.resize(img, (260, 260))  # 调整图像大小
            X.append(img.flatten())  # 将图像展平
            y.append(label_to_idx[label])

    return np.array(X), np.array(y)


# 加载数据集
data_dir = r'C:\Users\Harris\Documents\GitWarehouse\Dataset\d\CatDog'
X, y = load_dataset(data_dir)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 训练字典
start_train = time.time()
n_components = 100  # 字典中原子的数量
# dict_learner = MiniBatchDictionaryLearning(n_components=n_components, alpha=1, batch_size=10, transform_algorithm='omp')
dict_learner = DictionaryLearning(n_components=100)  # 设置字典中原子的数量
dict_learner.fit(X_train)
end_train = time.time()
train_time = end_train - start_train
print("Training time:", train_time, "s")

# 使用训练好的字典将训练集和测试集表示为稀疏编码
X_train_sparse = dict_learner.transform(X_train)
start_test = time.time()
X_test_sparse = dict_learner.transform(X_test)
end_test = time.time()
test_time = end_test - start_test
print("Testing time:", test_time, "s")




# 训练分类器（这里用简单的kNN分类器）
from sklearn.neighbors import KNeighborsClassifier
# 方案1 ：使用KNN分类器
classifier = KNeighborsClassifier(n_neighbors=5)
classifier.fit(X_train_sparse, y_train)

# # 方案2：使用随机森林分类器
# from sklearn.ensemble import RandomForestClassifier
#
# classifier = RandomForestClassifier(n_estimators=100)
# classifier.fit(X_train_sparse, y_train)

# # 方案3：使用SVM分类器
# from sklearn.svm import SVC
# classifier = SVC(kernel='linear')
# classifier.fit(X_train_sparse, y_train)

# 预测
y_pred = classifier.predict(X_test_sparse)

# 评估性能
accuracy = accuracy_score(y_test, y_pred)
print("Accuracy:", accuracy)
